from typing import Type, Callable, Tuple, List, Union

import torch
import torch.nn as nn
from timm.models.efficientnet_blocks import SqueezeExcite, DepthwiseSeparableConv
from timm.models.layers import drop_path, trunc_normal_, Mlp, DropPath


class MaxViT(nn.Module):
    """MaxViT模型的实现，来自论文：
    https://arxiv.org/pdf/2204.01697.pdf

    参数：
        in_channels (int, 可选): 输入通道数，默认值为3。
        depths (Tuple[int, ...], 可选): 每个网络阶段的深度，默认值为(2, 2, 5, 2)。
        channels (Tuple[int, ...], 可选): 每个网络阶段的通道数，默认值为(64, 128, 256, 512)。
        num_classes (int, 可选): 要预测的类别数，默认值为1000。
        embed_dim (int, 可选): 卷积stem的嵌入维度，默认值为64。
        num_heads (int, 可选): 注意力头的数量，默认值为32。
        grid_window_size (Tuple[int, int], 可选): 使用的网格/窗口大小，默认值为(7, 7)。
        attn_drop (float, 可选): 注意力权重的dropout比例，默认值为0.0。
        drop (float, 可选): 输出的dropout比例，默认值为0.0。
        drop_path (float, 可选): 路径的dropout比例，默认值为0.0。
        mlp_ratio (float, 可选): MLP隐藏层维度与嵌入维度的比例，默认值为4.0。
        act_layer (Type[nn.Module], 可选): 使用的激活函数类型，默认值为nn.GELU。
        norm_layer (Type[nn.Module], 可选): 使用的归一化层类型，默认值为nn.BatchNorm2d。
        norm_layer_transformer (Type[nn.Module], 可选): Transformer中的归一化层，默认值为nn.LayerNorm。
        debug (bool): 是否启用调试模式，默认值为False。
    """

    def __init__(
            self,
            in_channels: int = 3,
            depths: Tuple[int, ...] = (2, 2, 5, 2),
            channels: Tuple[int, ...] = (64, 128, 256, 512),
            embed_dim: int = 64,
            num_heads: int = 32,
            grid_window_size: Tuple[int, int] = (7, 7),
            attn_drop: float = 0.0,
            drop: float = 0.0,
            drop_path: float = 0.0,
            mlp_ratio: float = 4.0,
            act_layer: Type[nn.Module] = nn.GELU,
            norm_layer: Type[nn.Module] = nn.BatchNorm2d,
            norm_layer_transformer: Type[nn.Module] = nn.LayerNorm,
            debug: bool = False,
    ) -> None:
        """构造函数"""
        super(MaxViT, self).__init__()
        self.debug = debug
        # 检查参数
        assert len(depths) == len(channels), "每个阶段都必须指定通道数。"

        # 初始化卷积stem
        self.stem = nn.Sequential(
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=embed_dim,
                kernel_size=(3, 3),
                stride=(2, 2),
                padding=(1, 1),
            ),
            act_layer(),
            nn.Conv2d(
                in_channels=embed_dim,
                out_channels=embed_dim,
                kernel_size=(3, 3),
                stride=(1, 1),
                padding=(1, 1),
            ),
            act_layer(),
        )

        # 初始化各个阶段的模块
        # 生成一个线性增长的 drop path 概率序列
        drop_path_rates = torch.linspace(0.0, drop_path, sum(depths)).tolist()
        stages = []
        for index, (depth, channel) in enumerate(zip(depths, channels)):
            stages.append(
                MaxViTStage(
                    depth=depth,
                    in_channels=embed_dim if index == 0 else channels[index - 1],
                    out_channels=channel,
                    num_heads=num_heads,
                    grid_window_size=grid_window_size,
                    attn_drop=attn_drop,
                    drop=drop,
                    drop_path=drop_path_rates[sum(depths[:index]):sum(depths[:index + 1])],
                    mlp_ratio=mlp_ratio,
                    act_layer=act_layer,
                    norm_layer=norm_layer,
                    norm_layer_transformer=norm_layer_transformer,
                    debug=debug,
                )
            )
        self.stages = nn.ModuleList(stages)

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """前向传播

        参数：
            input (torch.Tensor): 输入图像，形状为[B, C, H, W]。

        返回：
            output (torch.Tensor): 分类输出，形状为[B, num_classes]。
        """
        outs = []
        x = self.stem(input)
        if self.debug:
            print('[MaxViT][forward]', '[stem][x]', x.shape)
        for i, stage in enumerate(self.stages):
            x = stage(x)
            outs.append(x)
            if self.debug:
                print('[MaxViT][forward_features]', f'[stage{i}][x]', x.shape)
        return outs


class MaxViTStage(nn.Module):
    """MaxViT的一个阶段。

    参数：
        depth (int): 阶段的深度。
        in_channels (int): 输入通道数。
        out_channels (int): 输出通道数。
        num_heads (int, 可选): 注意力头的数量，默认值为32。
        grid_window_size (Tuple[int, int], 可选): 使用的网格/窗口大小，默认值为(7, 7)。
        attn_drop (float, 可选): 注意力权重的dropout比例，默认值为0.0。
        drop (float, 可选): 输出的dropout比例，默认值为0.0。
        drop_path (float, 可选): 路径的dropout比例，默认值为0.0。
        mlp_ratio (float, 可选): MLP隐藏层维度与嵌入维度的比例，默认值为4.0。
        act_layer (Type[nn.Module], 可选): 使用的激活函数类型，默认值为nn.GELU。
        norm_layer (Type[nn.Module], 可选): 使用的归一化层类型，默认值为nn.BatchNorm2d。
        norm_layer_transformer (Type[nn.Module], 可选): Transformer中的归一化层，默认值为nn.LayerNorm。
        debug (bool): 是否启用调试模式，默认值为False。
    """

    def __init__(
            self,
            depth: int,
            in_channels: int,
            out_channels: int,
            num_heads: int = 32,
            grid_window_size: Tuple[int, int] = (7, 7),
            attn_drop: float = 0.0,
            drop: float = 0.0,
            drop_path: Union[List[float], float] = 0.0,
            mlp_ratio: float = 4.0,
            act_layer: Type[nn.Module] = nn.GELU,
            norm_layer: Type[nn.Module] = nn.BatchNorm2d,
            norm_layer_transformer: Type[nn.Module] = nn.LayerNorm,
            debug: bool = False,
    ) -> None:
        """构造函数"""
        super(MaxViTStage, self).__init__()
        self.module_name = '[MaxViTStage]'
        self.debug = debug
        if self.debug:
            print(self.module_name, '[depth]', depth)
            print(self.module_name, '[in_channels]', in_channels)
            print(self.module_name, '[out_channels]', out_channels)
            print(self.module_name, '[num_heads]', num_heads)
            print(self.module_name, '[grid_window_size]', grid_window_size)
            print(self.module_name, '[attn_drop]', attn_drop)
            print(self.module_name, '[drop]', drop)
            print(self.module_name, '[drop_path]', drop_path)
            print(self.module_name, '[mlp_ratio]', mlp_ratio)
        # 初始化块
        self.blocks = nn.Sequential(
            *[
                MaxViTBlock(
                    in_channels=in_channels if index == 0 else out_channels,
                    out_channels=out_channels,
                    downscale=index == 0,
                    num_heads=num_heads,
                    grid_window_size=grid_window_size,
                    attn_drop=attn_drop,
                    drop=drop,
                    drop_path=drop_path if isinstance(drop_path, float) else drop_path[index],
                    mlp_ratio=mlp_ratio,
                    act_layer=act_layer,
                    norm_layer=norm_layer,
                    norm_layer_transformer=norm_layer_transformer,
                    debug=debug,
                )
                for index in range(depth)
            ]
        )

    def forward(self, input=torch.Tensor) -> torch.Tensor:
        """前向传播

        参数：
            input (torch.Tensor): 输入张量，形状为[B, C_in, H, W]。

        返回：
            output (torch.Tensor): 输出张量，形状为[B, C_out, H // 2, W // 2]。
        """
        output = self.blocks(input)
        return output


class MaxViTBlock(nn.Module):
    """MaxViT块，由MBConv块、块级注意力和网格级注意力组成。

    参数：
        in_channels (int): 输入通道数。
        out_channels (int): 输出通道数。
        downscale (bool, 可选): 是否进行空间下采样，默认值为False。
        num_heads (int, 可选): 注意力头的数量，默认值为32。
        grid_window_size (Tuple[int, int], 可选): 使用的网格/窗口大小，默认值为(7, 7)。
        attn_drop (float, 可选): 注意力权重的dropout比例，默认值为0.0。
        drop (float, 可选): 输出的dropout比例，默认值为0.0。
        drop_path (float, 可选): 路径的dropout比例，默认值为0.0。
        mlp_ratio (float, 可选): MLP隐藏层维度与嵌入维度的比例，默认值为4.0。
        act_layer (Type[nn.Module], 可选): 使用的激活函数类型，默认值为nn.GELU。
        norm_layer (Type[nn.Module], 可选): 使用的归一化层类型，默认值为nn.BatchNorm2d。
        norm_layer_transformer (Type[nn.Module], 可选): Transformer中的归一化层，默认值为nn.LayerNorm。
        debug (bool): 是否启用调试模式，默认值为False。
    """

    def __init__(
            self,
            in_channels: int,
            out_channels: int,
            downscale: bool = False,
            num_heads: int = 32,
            grid_window_size: Tuple[int, int] = (7, 7),
            attn_drop: float = 0.0,
            drop: float = 0.0,
            drop_path: float = 0.0,
            mlp_ratio: float = 4.0,
            act_layer: Type[nn.Module] = nn.GELU,
            norm_layer: Type[nn.Module] = nn.BatchNorm2d,
            norm_layer_transformer: Type[nn.Module] = nn.LayerNorm,
            debug: bool = False,
    ) -> None:
        """构造函数"""
        super(MaxViTBlock, self).__init__()
        self.module_name = '[MaxViTBlock]'
        self.debug = debug
        if self.debug:
            print(self.module_name, '[in_channels]', in_channels)
            print(self.module_name, '[out_channels]', out_channels)
            print(self.module_name, '[downscale]', downscale)
            print(self.module_name, '[num_heads]', num_heads)
            print(self.module_name, '[grid_window_size]', grid_window_size)
            print(self.module_name, '[attn_drop]', attn_drop)
            print(self.module_name, '[drop]', drop)
            print(self.module_name, '[drop_path]', drop_path)
            print(self.module_name, '[mlp_ratio]', mlp_ratio)
        # 初始化MBConv块
        self.mb_conv = MBConv(
            in_channels=in_channels,
            out_channels=out_channels,
            downscale=downscale,
            act_layer=act_layer,
            norm_layer=norm_layer,
            drop_path=drop_path,
            debug=debug,
        )
        # 初始化块级和网格级Transformer
        self.block_transformer = MaxViTTransformerBlock(
            in_channels=out_channels,
            partition_function=window_partition,
            reverse_function=window_reverse,
            num_heads=num_heads,
            grid_window_size=grid_window_size,
            attn_drop=attn_drop,
            drop=drop,
            drop_path=drop_path,
            mlp_ratio=mlp_ratio,
            act_layer=act_layer,
            norm_layer=norm_layer_transformer,
            debug=debug,
        )
        self.grid_transformer = MaxViTTransformerBlock(
            in_channels=out_channels,
            partition_function=grid_partition,
            reverse_function=grid_reverse,
            num_heads=num_heads,
            grid_window_size=grid_window_size,
            attn_drop=attn_drop,
            drop=drop,
            drop_path=drop_path,
            mlp_ratio=mlp_ratio,
            act_layer=act_layer,
            norm_layer=norm_layer_transformer,
            debug=debug,
        )

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """前向传播

        参数：
            input (torch.Tensor): 输入张量，形状为[B, C_in, H, W]

        返回：
            output (torch.Tensor): 输出张量，形状为[B, C_out, H // 2, W // 2]（如果进行了下采样）
        """
        output = self.mb_conv(input)
        output = self.block_transformer(output)
        output = self.grid_transformer(output)
        return output


def window_partition(
        input: torch.Tensor,
        window_size: Tuple[int, int] = (7, 7)
) -> torch.Tensor:
    """窗口划分函数。

    参数：
        input (torch.Tensor): 输入张量，形状为[B, C, H, W]。
        window_size (Tuple[int, int], 可选): 要应用的窗口大小，默认值为(7, 7)。

    返回：
        windows (torch.Tensor): 展开的输入张量，形状为[B * 窗口数, window_size[0], window_size[1], C]。
    """
    # 获取输入的尺寸
    B, C, H, W = input.shape
    # 展开输入
    windows = input.view(
        B,
        C,
        H // window_size[0],
        window_size[0],
        W // window_size[1],
        window_size[1]
    )
    # 调整维度并重塑为[B * 窗口数, window_size[0], window_size[1], C]: (1,64,8,8,8,8) -> (1,8,8,8,8,64) -> (64,8,8,64)
    windows = windows.permute(0, 2, 4, 3, 5, 1).contiguous().view(
        -1, window_size[0], window_size[1], C
    )
    return windows


def window_reverse(
        windows: torch.Tensor,
        original_size: Tuple[int, int],
        window_size: Tuple[int, int] = (7, 7)
) -> torch.Tensor:
    """反转窗口划分操作。

    参数：
        windows (torch.Tensor): 窗口张量，形状为[B * 窗口数, window_size[0], window_size[1], C]。
        original_size (Tuple[int, int]): 原始尺寸 (H, W)。
        window_size (Tuple[int, int], 可选): 应用的窗口大小，默认值为(7, 7)。

    返回：
        output (torch.Tensor): 折叠后的输出张量，形状为[B, C, H, W]。
    """
    # 获取高度和宽度
    H, W = original_size
    # 计算原始的批量大小
    B = int(windows.shape[0] / (H * W / window_size[0] / window_size[1]))
    # 折叠张量
    output = windows.view(
        B,
        H // window_size[0],
        W // window_size[1],
        window_size[0],
        window_size[1],
        -1
    )
    output = output.permute(0, 5, 1, 3, 2, 4).contiguous().view(
        B, -1, H, W
    )
    return output


def grid_partition(
        input: torch.Tensor,
        grid_size: Tuple[int, int] = (7, 7)
) -> torch.Tensor:
    """网格划分函数。

    参数：
        input (torch.Tensor): 输入张量，形状为[B, C, H, W]。
        grid_size (Tuple[int, int], 可选): 要应用的网格大小，默认值为(7, 7)。

    返回：
        grid (torch.Tensor): 展开的输入张量，形状为[B * 网格数, grid_size[0], grid_size[1], C]。
    """
    # 获取输入的尺寸
    B, C, H, W = input.shape
    # 展开输入
    grid = input.view(
        B,
        C,
        grid_size[0],
        H // grid_size[0],
        grid_size[1],
        W // grid_size[1]
    )
    # 调整维度并重塑为[B * (H // grid_size[0]) * (W // grid_size[1]), grid_size[0], grid_size[1], C]
    grid = grid.permute(0, 3, 5, 2, 4, 1).contiguous().view(
        -1, grid_size[0], grid_size[1], C
    )
    return grid


def grid_reverse(
        grid: torch.Tensor,
        original_size: Tuple[int, int],
        grid_size: Tuple[int, int] = (7, 7)
) -> torch.Tensor:
    """反转网格划分操作。

    参数：
        grid (torch.Tensor): 网格张量，形状为[B * 网格数, grid_size[0], grid_size[1], C]。
        original_size (Tuple[int, int]): 原始尺寸 (H, W)。
        grid_size (Tuple[int, int], 可选): 应用的网格大小，默认值为(7, 7)。

    返回：
        output (torch.Tensor): 折叠后的输出张量，形状为[B, C, H, W]。
    """
    # 获取高度、宽度和通道数
    H, W = original_size
    C = grid.shape[-1]
    # 计算原始的批量大小
    B = int(grid.shape[0] / (H * W / grid_size[0] / grid_size[1]))
    # 折叠张量
    output = grid.view(
        B,
        H // grid_size[0],
        W // grid_size[1],
        grid_size[0],
        grid_size[1],
        C
    )
    output = output.permute(0, 5, 3, 1, 4, 2).contiguous().view(
        B, C, H, W
    )
    return output


def get_relative_position_index(
        win_h: int,
        win_w: int
) -> torch.Tensor:
    """生成窗口内每个token的两两相对位置索引的函数。
    参考自Timm的Swin V1实现。

    参数：
        win_h (int): 窗口/网格的高度。
        win_w (int): 窗口/网格的宽度。

    返回：
        relative_coords (torch.Tensor): 两两相对位置索引，形状为[height * width, height * width]。
    """
    coords = torch.stack(
        torch.meshgrid([torch.arange(win_h), torch.arange(win_w)], indexing='ij')
    )
    coords_flatten = torch.flatten(coords, 1)
    relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
    relative_coords = relative_coords.permute(1, 2, 0).contiguous()
    relative_coords[:, :, 0] += win_h - 1
    relative_coords[:, :, 1] += win_w - 1
    relative_coords[:, :, 0] *= 2 * win_w - 1
    return relative_coords.sum(-1)


class MaxViTTransformerBlock(nn.Module):
    """MaxViT Transformer块。

    使用块划分：
    x ← x + Unblock(RelAttention(Block(LN(x))))
    x ← x + MLP(LN(x))

    使用网格划分：
    x ← x + Ungrid(RelAttention(Grid(LN(x))))
    x ← x + MLP(LN(x))

    为了避免多次重塑操作，层归一化（LN）在网格/窗口划分后应用。
    最终输出时执行网格/窗口反转（Unblock/Ungrid），原因相同。

    参数：
        in_channels (int): 输入通道数。
        partition_function (Callable): 要使用的划分函数（网格或窗口划分）。
        reverse_function (Callable): 要使用的反转函数（网格或窗口反转）。
        num_heads (int, 可选): 注意力头的数量，默认值为32。
        grid_window_size (Tuple[int, int], 可选): 使用的网格/窗口大小，默认值为(7, 7)。
        attn_drop (float, 可选): 注意力权重的dropout比例，默认值为0.0。
        drop (float, 可选): 输出的dropout比例，默认值为0.0。
        drop_path (float, 可选): 路径的dropout比例，默认值为0.0。
        mlp_ratio (float, 可选): MLP隐藏层维度与嵌入维度的比例，默认值为4.0。
        act_layer (Type[nn.Module], 可选): 使用的激活函数类型，默认值为nn.GELU。
        norm_layer (Type[nn.Module], 可选): 使用的归一化层类型，默认值为nn.LayerNorm。
        debug (bool): 是否启用调试模式，默认值为False。
    """

    def __init__(
            self,
            in_channels: int,
            partition_function: Callable,
            reverse_function: Callable,
            num_heads: int = 32,
            grid_window_size: Tuple[int, int] = (7, 7),
            attn_drop: float = 0.0,
            drop: float = 0.0,
            drop_path: float = 0.0,
            mlp_ratio: float = 4.0,
            act_layer: Type[nn.Module] = nn.GELU,
            norm_layer: Type[nn.Module] = nn.LayerNorm,
            debug: bool = False
    ) -> None:
        """构造函数"""
        super(MaxViTTransformerBlock, self).__init__()
        self.module_name = '[MaxViTTransformerBlock]'
        self.debug = debug
        # 保存参数
        self.partition_function = partition_function
        self.reverse_function = reverse_function
        self.grid_window_size = grid_window_size
        # 初始化层
        self.norm_1 = norm_layer(in_channels)
        self.attention = RelativeSelfAttention(
            in_channels=in_channels,
            num_heads=num_heads,
            grid_window_size=grid_window_size,
            attn_drop=attn_drop,
            drop=drop
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm_2 = norm_layer(in_channels)
        self.mlp = Mlp(
            in_features=in_channels,
            hidden_features=int(mlp_ratio * in_channels),
            act_layer=act_layer,
            drop=drop
        )

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """前向传播。

        参数：
            input (torch.Tensor): 输入张量，形状为[B, C_in, H, W]。

        返回：
            output (torch.Tensor): 输出张量，形状为[B, C_out, H, W]。
        """
        # 保存原始形状
        B, C, H, W = input.shape
        # 执行划分
        input_partitioned = self.partition_function(input, self.grid_window_size)
        input_partitioned = input_partitioned.view(
            -1, self.grid_window_size[0] * self.grid_window_size[1], C
        )
        # 执行归一化、注意力和dropout
        output = input_partitioned + self.drop_path(
            self.attention(self.norm_1(input_partitioned))
        )
        # 执行归一化、MLP和dropout
        output = output + self.drop_path(self.mlp(self.norm_2(output)))
        # 反转划分
        output = self.reverse_function(output, (H, W), self.grid_window_size)
        return output


class MBConv(nn.Module):
    """MBConv块，描述于：https://arxiv.org/pdf/2204.01697.pdf。

    无下采样：
    x ← x + Proj(SE(DWConv(Conv(Norm(x)))))

    有下采样：
    x ← Proj(Pool2D(x)) + Proj(SE(DWConv ↓(Conv(Norm(x))))).

    Conv是一个1x1卷积，后接批量归一化层和GELU激活函数。
    SE是Squeeze-Excitation层。
    Proj是收缩的1x1卷积。

    注意：此实现与原始的MobileNet实现略有不同！

    参数：
        in_channels (int): 输入通道数。
        out_channels (int): 输出通道数。
        downscale (bool, 可选): 是否执行下采样（缩小一半尺寸），默认值为False。
        act_layer (Type[nn.Module], 可选): 使用的激活函数类型，默认值为nn.GELU。
        norm_layer (Type[nn.Module], 可选): 使用的归一化层类型，默认值为nn.BatchNorm2d。
        drop_path (float, 可选): 训练期间应用的dropout率，默认值为0.0。
        debug (bool): 是否启用调试模式，默认值为False。
    """

    def __init__(
            self,
            in_channels: int,
            out_channels: int,
            downscale: bool = False,
            act_layer: Type[nn.Module] = nn.GELU,
            norm_layer: Type[nn.Module] = nn.BatchNorm2d,
            drop_path: float = 0.0,
            debug: bool = False
    ) -> None:
        """构造函数"""
        super(MBConv, self).__init__()
        self.module_name = '[MBConv]'
        self.debug = debug
        # 保存参数
        self.drop_path_rate = drop_path
        # 检查下采样参数
        if not downscale:
            assert in_channels == out_channels, "如果不进行下采样，输入和输出通道数必须相等。"
        # 如果使用GELU，则忽略inplace参数
        if act_layer == nn.GELU:
            act_layer = _gelu_ignore_parameters
        # 构建主路径
        self.main_path = nn.Sequential(
            norm_layer(in_channels),
            nn.Conv2d(in_channels, in_channels, kernel_size=1),
            DepthwiseSeparableConv(
                in_chs=in_channels,
                out_chs=out_channels,
                stride=2 if downscale else 1,  # 通过此进行下采样
                act_layer=act_layer,
                norm_layer=norm_layer,
                drop_path_rate=drop_path
            ),
            SqueezeExcite(in_chs=out_channels, rd_ratio=0.25),
            nn.Conv2d(out_channels, out_channels, kernel_size=1)
        )
        # 构建跳跃路径
        self.skip_path = nn.Sequential(
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels, out_channels, kernel_size=1)
        ) if downscale else nn.Identity()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """前向传播。

        参数：
            input (torch.Tensor): 输入张量，形状为[B, C_in, H, W]。

        返回：
            output (torch.Tensor): 输出张量，形状为[B, C_out, H // 2, W // 2]（如果进行了下采样）。
        """
        output = self.main_path(input)
        if self.drop_path_rate > 0.0:
            output = drop_path(output, self.drop_path_rate, self.training)
        output = output + self.skip_path(input)
        return output


class RelativeSelfAttention(nn.Module):
    """类似于Swin V1的相对自注意力机制。实现参考Timm的Swin V1实现。

    参数：
        in_channels (int): 输入通道数。
        num_heads (int, 可选): 注意力头的数量，默认值为32。
        grid_window_size (Tuple[int, int], 可选): 使用的网格/窗口大小，默认值为(7, 7)。
        attn_drop (float, 可选): 注意力权重的dropout比例，默认值为0.0。
        drop (float, 可选): 输出的dropout比例，默认值为0.0。
    """

    def __init__(
            self,
            in_channels: int,
            num_heads: int = 32,
            grid_window_size: Tuple[int, int] = (7, 7),
            attn_drop: float = 0.0,
            drop: float = 0.0
    ) -> None:
        """构造函数"""
        super(RelativeSelfAttention, self).__init__()
        # 保存参数
        self.in_channels = in_channels
        self.num_heads = num_heads
        self.grid_window_size = grid_window_size
        self.scale = num_heads ** -0.5
        self.attn_area = grid_window_size[0] * grid_window_size[1]
        # 初始化层
        self.qkv_mapping = nn.Linear(in_channels, 3 * in_channels, bias=True)
        self.attn_drop = nn.Dropout(p=attn_drop)
        self.proj = nn.Linear(in_channels, in_channels, bias=True)
        self.proj_drop = nn.Dropout(p=drop)
        self.softmax = nn.Softmax(dim=-1)
        # 定义相对位置偏差的参数表，形状为[2*Wh-1 * 2*Ww-1, nH]
        self.relative_position_bias_table = nn.Parameter(
            torch.zeros(
                (2 * grid_window_size[0] - 1) * (2 * grid_window_size[1] - 1),
                num_heads
            )
        )
        # 获取窗口内每个token的两两相对位置索引
        self.register_buffer(
            "relative_position_index",
            get_relative_position_index(grid_window_size[0], grid_window_size[1])
        )
        # 初始化相对位置偏差
        trunc_normal_(self.relative_position_bias_table, std=0.02)

    def _get_relative_positional_bias(self) -> torch.Tensor:
        """返回相对位置偏差。

        返回：
            relative_position_bias (torch.Tensor): 相对位置偏差张量。
        """
        relative_position_bias = self.relative_position_bias_table[
            self.relative_position_index.view(-1)
        ].view(self.attn_area, self.attn_area, -1)
        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
        return relative_position_bias.unsqueeze(0)

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """前向传播。

        参数：
            input (torch.Tensor): 输入张量，形状为[B_, N, C]。

        返回：
            output (torch.Tensor): 输出张量，形状为[B_, N, C]。
        """
        # 获取输入的形状
        B_, N, C = input.shape
        # 执行QKV映射
        qkv = self.qkv_mapping(input).reshape(
            B_, N, 3, self.num_heads, -1
        ).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        # 缩放查询向量
        q = q * self.scale
        # 计算注意力图
        attn = self.softmax(
            q @ k.transpose(-2, -1) + self._get_relative_positional_bias()
        )
        # 使用注意力图映射值向量
        output = (attn @ v).transpose(1, 2).reshape(B_, N, -1)
        # 执行最终的投影和dropout
        output = self.proj(output)
        output = self.proj_drop(output)
        return output


def _gelu_ignore_parameters(*args, **kwargs) -> nn.Module:
    """忽略DepthwiseSeparableConv中inplace=True参数的技巧。

    参数：
        *args: 忽略。
        **kwargs: 忽略。

    返回：
        activation (nn.Module): GELU激活函数。
    """
    activation = nn.GELU()
    return activation


if __name__ == '__main__':
    model = MaxViT(
        in_channels=3,
        depths=(2, 2, 5, 2),
        channels=(64, 128, 256, 512),
        embed_dim=64,
        grid_window_size=(8, 8),
        attn_drop=0.2,
        drop=0.2,
        drop_path=0.2,
        debug=False,
    )
    input = torch.rand(1, 3, 256, 256)
    output = model(input)
    print(f'Total parameters: {sum(param.numel() for param in model.parameters()) / 1e6:.2f} M')
    print(output[3].shape)
